package simpletrace

import (
	"context"
	"fmt"
	"strconv"
)


type RedisSpanReader struct {
	redisClient RedisClient // Assume RedisClient is defined elsewhere
}

func NewRedisSpanReader(redisClient RedisClient) *RedisSpanReader {
	return &RedisSpanReader{redisClient: redisClient}
}

func (r *RedisSpanReader) Read(ctx context.Context, traceID string) (*Span, error) {
	// 从reids中获取所有span，并组装成一棵完整的树返回
	key := fmt.Sprintf(RedisKeyTraceSpanIDSet, traceID)
	spanIDs, err := r.redisClient.SMembersWithCtx(ctx, key)
	if err != nil {
		return nil, fmt.Errorf("[strace] failed to get span IDs for trace %s: %w", traceID, err)
	}
	if len(spanIDs) == 0 {
		return nil, fmt.Errorf("[strace] no spans found for trace %s", traceID)
	}
	spans := make([]*Span, 0, len(spanIDs))
	for _, spanID := range spanIDs {
		spanKey := fmt.Sprintf(RedisKeyTraceSpanInfo, traceID, spanID)
		data, err := r.redisClient.Bytes(r.redisClient.GetWithCtx(ctx, spanKey))
		if err != nil {
			return nil, fmt.Errorf("[strace] failed to get span %s for trace %s: %w", spanID, traceID, err)
		}
		if data == nil {
			continue // Skip if no data found for this span
		}
		span := &Span{}
		if err := span.Unmarshal(data); err != nil {
			return nil, fmt.Errorf("[strace] failed to unmarshal span %s for trace %s: %w", spanID, traceID, err)
		}
		spans = append(spans, span)
	}
	logger.WarnCtx(ctx, "[strace] read trace %s with %d spans", traceID, len(spans))
	return buildTraceTree(spans), nil
}

func (r *RedisSpanReader) GetAllSpanIDs(ctx context.Context, traceID string) ([]string, error) {
	key := fmt.Sprintf(RedisKeyTraceSpanIDSet, traceID)
	spanIDs, err := r.redisClient.SMembersWithCtx(ctx, key)
	if err != nil {
		return nil, fmt.Errorf("failed to get span IDs for trace %s: %w", traceID, err)
	}
	return spanIDs, nil
}

func (r *RedisSpanReader) GetSpanDataStr(ctx context.Context, traceID, spanID string) (string, error) {
	key := fmt.Sprintf(RedisKeyTraceSpanInfo, traceID, spanID)
	data, err := r.redisClient.Bytes(r.redisClient.GetWithCtx(ctx, key))
	if err != nil {
		return "", fmt.Errorf("failed to get span %s for trace %s: %w", spanID, traceID, err)
	}
	if data == nil {
		return "", fmt.Errorf("no data found for span %s in trace %s", spanID, traceID)
	}
	return string(data), nil
}

func (r *RedisSpanReader) GetTraceRootSpans(ctx context.Context, traceIDs []string) ([]*Span, error) {
	if len(traceIDs) == 0 {
		return nil, fmt.Errorf("[strace] no trace IDs provided")
	}

	spans := make([]*Span, 0, len(traceIDs))
	keys := make([]string, 0, len(traceIDs))
	for _, traceID := range traceIDs {
		keys = append(keys, fmt.Sprintf(RedisKeyTraceRootSpanID, traceID))
	}
	results, err := r.redisClient.MGetWithCtx(ctx, keys...)
	if err != nil {
		return nil, fmt.Errorf("[strace] failed to get root spans for trace IDs %v: %w", traceIDs, err)
	}
	for _, data := range results {
		if data == nil {
			continue // Skip if no data found for this trace ID
		}
		span := &Span{}
		if err := span.Unmarshal([]byte(data.(string))); err != nil {
			logger.ErrorCtx(ctx, "[strace] failed to unmarshal root span %v", err)
			continue // Skip this span if unmarshalling fails
		}
		spans = append(spans, span)
	}
	return spans, nil
}

func (r *RedisSpanReader) GetAllTraces(ctx context.Context) ([]string, error) {
	keys, err := r.redisClient.SMembersWithCtx(ctx, RedisKeyTraceIDSet)
	if err != nil {
		return nil, fmt.Errorf("[strace] failed to get all trace keys: %w", err)
	}
	traceIDs := make([]string, 0, len(keys))
	for _, key := range keys {
		if key == "" {
			continue // Skip empty keys
		}
		traceIDs = append(traceIDs, key)
	}
	return traceIDs, nil
}

func (r *RedisSpanReader) CountSpans(ctx context.Context, traceID string) (int64, error) {
	key := fmt.Sprintf(RedisKeyTraceSpanIDSet, traceID)
	count, err := r.redisClient.SCardWithCtx(ctx, key)
	if err != nil {
		return 0, fmt.Errorf("[strace] failed to count spans for trace %s: %w", traceID, err)
	}
	return count, nil
}

func (r *RedisSpanReader) Clean(ctx context.Context) error {
	traceIDs, err := r.GetAllTraces(ctx)
	if err != nil {
		return fmt.Errorf("[strace] failed to get all trace IDs for cleanup: %w", err)
	}
	if len(traceIDs) == 0 {
		logger.WarnCtx(ctx, "[strace] No traces to clean")
		return nil // Nothing to clean
	}
	// 删除所有trace相关的key
	for _, traceID := range traceIDs {
		r.CleanByTraceID(ctx, traceID)
	}
	// 删除trace ID集合
	if _, err := r.redisClient.DelWithCtx(ctx, RedisKeyTraceIDSet); err != nil {
		logger.ErrorCtx(ctx, "[strace] failed to delete trace ID set: %v", err)
		return err
	}
	return nil
}

func (r *RedisSpanReader) CleanByTraceID(ctx context.Context, traceID string) error {
	logger.WarnCtx(ctx, "[strace] cleaning trace %s", traceID)
	// 获取trace的span ID集合和根span ID
	spanIDs, err := r.GetAllSpanIDs(ctx, traceID)
	if err != nil {
		logger.ErrorCtx(ctx, "[strace] failed to get span IDs for trace %s: %v", traceID, err)
	}
	spanSetKey := fmt.Sprintf(RedisKeyTraceSpanIDSet, traceID)
	if _, err := r.redisClient.DelWithCtx(ctx, spanSetKey); err != nil {
		logger.ErrorCtx(ctx, "[strace] failed to delete keys for trace %s: %v", traceID, err)
	}
	rootSpanKey := fmt.Sprintf(RedisKeyTraceRootSpanID, traceID)
	if _, err := r.redisClient.DelWithCtx(ctx, rootSpanKey); err != nil {
		logger.ErrorCtx(ctx, "[strace] failed to delete keys for trace %s: %v", traceID, err)
	}
	// 删除每个span的详细信息
	for _, spanID := range spanIDs {
		spanKey := fmt.Sprintf(RedisKeyTraceSpanInfo, traceID, spanID)
		if _, err := r.redisClient.DelWithCtx(ctx, spanKey); err != nil {
			logger.ErrorCtx(ctx, "[strace] failed to delete span %s for trace %s: %v", spanID, traceID, err)
		}
	}
	return nil
}

func (r *RedisSpanReader) GetTraceMemUsage(ctx context.Context, traceID string) (int64, error) {
	key := fmt.Sprintf(RedusKeyTraceMemUsage, traceID)
	memUsageStr, err := r.redisClient.String(r.redisClient.GetWithCtx(ctx, key))
	if err != nil {
		return 0, fmt.Errorf("[strace] failed to get memory usage for trace %s: %w", traceID, err)
	}
	memUsage, _ := strconv.ParseInt(memUsageStr, 10, 64)
	return memUsage, nil
}